from numpy import (
    ndarray,  # 导入ndarray类
    sqrt,  # 导入平方根函数
    matmul,  # 导入矩阵乘法函数
    trace,  # 导入矩阵迹函数
)

def norm_precoder(
        precoding_matrix: ndarray,  # 预编码矩阵，维度为 (sat_nr * ant_nr, user_nr)
        power_constraint_watt: float,  # 功率约束，单位为瓦特（W）
        per_satellite: bool,  # 是否对每颗卫星单独归一化
        sat_nr: int = 1,  # 卫星数量，默认为1
        sat_ant_nr: int = 1  # 每颗卫星的天线数量，默认为1
) -> ndarray:  # 返回归一化后的预编码矩阵
    """
    归一化预编码矩阵，使其满足功率约束

    参数:
        precoding_matrix (ndarray): 预编码矩阵，维度为 (sat_nr * ant_nr, user_nr)，
                                    其中包含所有卫星和天线的预编码矩阵。
        power_constraint_watt (float): 功率约束，单位为瓦特（W）。
        per_satellite (bool): 是否对每颗卫星单独进行归一化。如果为 True，则对每颗卫星单独归一化。
        sat_nr (int, optional): 卫星数量，默认为1。
        sat_ant_nr (int, optional): 每颗卫星的天线数量，默认为1。

    返回:
        ndarray: 归一化后的预编码矩阵，维度与输入矩阵相同。
    """

    if per_satellite:
        # 如果需要对每颗卫星单独归一化
        for satellite_id in range(sat_nr):
            # 获取当前卫星的预编码矩阵片段
            w_precoder_slice = precoding_matrix[satellite_id * sat_ant_nr: satellite_id * sat_ant_nr + sat_ant_nr, :]
            
            # 计算当前片段的归一化因子
            norm_factor_slice = sqrt(
                power_constraint_watt / sat_nr / trace(matmul(w_precoder_slice.conj().T, w_precoder_slice))
            )
            
            # 归一化当前片段的预编码矩阵
            w_precoder_slice_normed = norm_factor_slice * w_precoder_slice
            
            # 将归一化后的片段更新到原始矩阵中
            precoding_matrix[satellite_id * sat_ant_nr: satellite_id * sat_ant_nr + sat_ant_nr, :] = w_precoder_slice_normed

        # 返回归一化后的预编码矩阵
        normalized_precoder = precoding_matrix

    else:
        # 如果不需要对每颗卫星单独归一化
        # 计算整体的归一化因子
        norm_factor = sqrt(power_constraint_watt / trace(matmul(precoding_matrix.conj().T, precoding_matrix)))
        
        # 归一化整个预编码矩阵
        normalized_precoder = norm_factor * precoding_matrix

    return normalized_precoder
